package com.fa4j.common.base.util;


import com.fa4j.common.base.model.Tree;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树结构数据转化
 */
public class TreeUtil {

    public static <ID, T extends Tree<ID, T>> List<T> listToTree(List<T> list, ID rootId) {
        return listToTree(list, rootId, Tree::getId, Tree::getParentId, Tree::setChildren, Tree::getPriority);
    }

    public static <ID, T extends Tree<ID, T>> List<T> listToTree(List<T> dataList, ID rootId, Function<T, ID> idFun, Function<T, ID> parentIdFun, BiConsumer<T, List<T>> setChildrenFun, Function<T, Integer> keyExtractor) {
        if (dataList == null || dataList.isEmpty()) return dataList;
        Map<ID, List<T>> parentIdListMap = dataList.stream()
                .sorted(Comparator.comparing(keyExtractor))
                .collect(Collectors.groupingBy(parentIdFun));
        List<T> rootList = parentIdListMap.get(rootId);
        for (T data : rootList) {
            setChildren(data, parentIdListMap, idFun, setChildrenFun);
        }
        return rootList;
    }

    private static <T, ID> void setChildren(T data, Map<ID, List<T>> parentIdListMap, Function<T, ID> idFun, BiConsumer<T, List<T>> setChildrenFun) {
        if (data == null) return;
        List<T> children = parentIdListMap.get(idFun.apply(data));
        if (children == null || children.isEmpty()) return;

        setChildrenFun.accept(data, children);
        for (T child : children) {
            setChildren(child, parentIdListMap, idFun, setChildrenFun);
        }
    }
}
